// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#pragma once

#include <cantor/proto/parser/ast.pb.h>
#include <turbo/container/flat_hash_map.h>
#include <turbo/utility/status.h>
#include <turbo/log/logging.h>
#include <memory>
#include <variant>
#include <cantor/unresolved/stmt_traits.h>

namespace cantor::unresolved {
    /// none stmt
    ///
    struct DummyStmt {
    };



    using StmtVisitor = std::variant<
        DummyStmt,
        UnresolvedSelectStmt,
        UnresolvedStartTxnStmt,
        UnresolvedCommitTxnStmt,
        UnresolvedRollbackTxnStmt,
        UnresolvedKillStmt,
        UnresolvedTypeStmt
    >;
    ///////////////////////////////////////////////////////////////////////////////
    /// @brief UnresolvedStatement is the lightweight wrapper for any unresolved
    ///        SQL statement parsed from protobuf.
    ///
    /// This class acts as a generic container for a `cantor::pb::ast::Statement*`,
    /// providing runtime inspection, type-checking, and visitor-based casting.
    ///
    /// It is designed to:
    /// - Avoid slicing by storing a pointer to the actual statement.
    /// - Delay instantiation of specific typed AST nodes until explicitly cast.
    /// - Provide CRTP-based visitor access to the underlying statement.
    /// - Ensure safety and clarity with explicit cast, type checking, and error handling.
    ///
    /// You should **retain this object** as long as the visitor object needs to access
    /// the underlying `Statement`. All visitor instances reference memory inside this.
    ///
    /// Example usage:
    /// ```cpp
    /// std::vector<UnresolvedStatement> stmts;
    /// ...
    /// if (stmts[i].is_type<UnresolvedStartTxnStmt>()) {
    ///     auto txn = stmts[i].cast_or_die<UnresolvedStartTxnStmt>();
    ///     if (txn.is_read_only()) { ... }
    /// }
    /// ```
    ///////////////////////////////////////////////////////////////////////////////


    class UnresolvedStatement {
    public:
        static constexpr const cantor::pb::ast::StmtType TYPE = cantor::pb::ast::StmtType::STMT_NONE;

    public:
        UnresolvedStatement() = default;

        virtual ~UnresolvedStatement() = default;

        explicit UnresolvedStatement(const cantor::pb::ast::Statement *st) : stmt(st) {
        }


        void attach(const cantor::pb::ast::Statement *st) {
            stmt = st;
        }

        const cantor::pb::ast::Statement *stmt{nullptr};

    public:
        /// @brief Cast the underlying statement to a concrete typed visitor.
        ///
        /// This verifies that the `stmt->type()` matches `Visitor::TYPE`, and then
        /// constructs the visitor.
        ///
        /// @tparam Visitor A class derived from `UnresolvedVisitor<Visitor>`
        /// @return `turbo::Result<Visitor>` if successful, otherwise contains error.
        template<class Visitor>
        turbo::Result<Visitor> cast_to() const {
            if (!stmt) {
                return turbo::data_loss_error("no stmt setting");
            }
            if (stmt->type() != Visitor::TYPE && Visitor::TYPE != cantor::pb::ast::StmtType::STMT_NONE) {
                return turbo::invalid_argument_error("Failed to cast statement to type - statement type mismatch [",
                                                     cantor::pb::ast::StmtType_Name(stmt->type()), ":",
                                                     cantor::pb::ast::StmtType_Name(Visitor::TYPE), "]");
            }
            return Visitor(*stmt);
        }

        /// @brief Check if the stored statement matches the `Visitor::TYPE`.
        template<class Visitor>
        bool is_type() const {
            return stmt && (stmt->type() == Visitor::TYPE || Visitor::TYPE == cantor::pb::ast::StmtType::STMT_NONE);
        }

        /// @brief Check if the stored statement matches the specified statement type.
        template<cantor::pb::ast::StmtType Stype>
        bool is_stmt() const {
            return stmt && (stmt->type() == Stype || Stype == cantor::pb::ast::StmtType::STMT_NONE);
        }

        /// @brief Cast the statement to a visitor or crash fatally on error.
        ///
        /// Use this only when you're confident about the statement type.
        /// Crashes the program with a fatal log on mismatch or null.
        template<class Visitor>
        Visitor cast_or_die() const {
            auto r = cast_to<Visitor>();
            if (!r.ok()) {
                KLOG(FATAL) << r.status();
            }
            return std::move(r).value_or_die();
        }

        StmtVisitor cast_to(cantor::pb::ast::StmtType Stype) const {
            switch (Stype) {
                case cantor::pb::ast::StmtType::STMT_SELECT:
                    return cast_or_die<StmtTypeTrait<cantor::pb::ast::StmtType::STMT_SELECT>::type>();
                case cantor::pb::ast::StmtType::STMT_START_TRANSACTION:
                    return cast_or_die<StmtTypeTrait<cantor::pb::ast::StmtType::STMT_START_TRANSACTION>::type>();
                case cantor::pb::ast::StmtType::STMT_COMMIT_TRANSACTION:
                    return cast_or_die<StmtTypeTrait<cantor::pb::ast::StmtType::STMT_COMMIT_TRANSACTION>::type>();
                case cantor::pb::ast::StmtType::STMT_ROLLBACK_TRANSACTION:
                    return cast_or_die<StmtTypeTrait<cantor::pb::ast::StmtType::STMT_ROLLBACK_TRANSACTION>::type>();
                case cantor::pb::ast::StmtType::STMT_KILL:
                    return cast_or_die<StmtTypeTrait<cantor::pb::ast::StmtType::STMT_KILL>::type>();
                case cantor::pb::ast::StmtType::STMT_TYPE_NAME:
                    return cast_or_die<StmtTypeTrait<cantor::pb::ast::StmtType::STMT_TYPE_NAME>::type>();
                default:
                    return DummyStmt{};
            }
        }

        StmtVisitor cast_to() const {
            return cast_to(stmt->type());
        }
    };
} // namespace cantor::unresolved
